#include "Octree.h"
#include "Render//VulkanRHI.h"
#include "Utils/FileUtil.h"
#include "Core/CoreGlobals.h"
#include "Templates/NewObject.h"

Octree::Octree(glm::vec3 c1, glm::vec3 c2, int32 depth)
{
	mAA = c1;
	mBB = c2;
	mCenter = (c1 + c2) * 0.5f;
	mDepth = depth;
	mBallNum = 0;
	mHasChildren = false;
}

Octree::~Octree()
{
	if (mHasChildren)
	{
		DestroyChildren();
	}
}

void Octree::Add(Ball* ball)
{
	mBallNum++;
	if (!mHasChildren && mDepth < MAX_OCTREE_DEPTH && mBallNum > MAX_BALLS_PER_OCTREE)
	{
		CreateChildren();
	}

	if (mHasChildren)
	{
		FileBall(ball, ball->position, true);
	}
	else
	{
		mBalls.insert(ball);
	}
}

void Octree::Remove(Ball* ball)
{
	Remove(ball, ball->position);
}

void Octree::MoveBall(Ball* ball, glm::vec3 oldPos)
{
	Remove(ball, oldPos);
	Add(ball);
}

void Octree::FileBall(Ball* ball, glm::vec3 pos, bool isAddBall)
{
	for (int x = 0; x < 2; x++)
	{
		if (x == 0)
		{
			if (pos[0] - ball->radius > mCenter[0])
			{
				continue;
			}
		}
		else if (pos[0] + ball->radius < mCenter[0])
		{
			continue;
		}

		for (int y = 0; y < 2; y++)
		{
			if (y == 0)
			{
				if (pos[1] - ball->radius > mCenter[1])
				{
					continue;
				}
			}
			else if (pos[1] + ball->radius < mCenter[1])
			{
				continue;
			}

			for (int z = 0; z < 2; z++)
			{
				if (z == 0)
				{
					if (pos[2] - ball->radius > mCenter[2])
					{
						continue;
					}
				}
				else if (pos[2] + ball->radius < mCenter[2])
				{
					continue;
				}

				if (isAddBall)
				{
					mChildren[x][y][z]->Add(ball);
				}
				else
				{
					mChildren[x][y][z]->Remove(ball, pos);
				}
			}
		}
	}
}


void Octree::CreateChildren()
{
	for (int x = 0; x < 2; x++)
	{
		float minX, maxX;
		if (x == 0)
		{
			minX = mAA[0];
			maxX = mCenter[0];
		}
		else
		{
			minX = mCenter[0];
			maxX = mBB[0];
		}

		for (int y = 0; y < 2; y++)
		{
			float minY = 0.f;
			float maxY = 0.f;
			if (y == 0)
			{
				minX = mAA[1];
				maxX = mCenter[1];
			}
			else
			{
				minX = mCenter[1];
				maxX = mBB[1];
			}

			for (int z = 0; z < 2; z++)
			{
				float minZ = 0.f;
				float maxZ = 0.f;
				if (z == 0)
				{
					minX = mAA[2];
					maxX = mCenter[2];
				}
				else
				{
					minX = mCenter[2];
					maxX = mBB[2];
				}

				mChildren[x][y][z] = NewObject<Octree>(glm::vec3(minX, minY, minZ),
					glm::vec3(maxX, maxY, maxZ),
					mDepth + 1);
			}
		}
	}

	for (std::set<Ball*>::iterator it = mBalls.begin(); it != mBalls.end(); it++)
	{
		Ball* ball = *it;
		FileBall(ball, ball->position, true);
	}
	mBalls.clear();
	mHasChildren = true;
}

void Octree::CollectBalls(std::set<Ball*>& bs)
{
	if (mHasChildren)
	{
		for (int x = 0; x < 2; x++)
		{
			for (int y = 0; y < 2; y++)
			{
				for (int z = 0; z < 2; z++)
				{
					mChildren[x][y][z]->CollectBalls(bs);
				}
			}
		}
	}
	else
	{
		for (std::set<Ball*>::iterator it = mBalls.begin(); it != mBalls.end(); it++)
		{
			bs.insert(*it);
		}
	}
}

void Octree::DestroyChildren()
{
	CollectBalls(mBalls);

	for (int x = 0; x < 2; x++)
	{
		for (int y = 0; y < 2; y++)
		{
			for (int z = 0; z < 2; z++)
			{
				DeleteObject<Octree>(mChildren[x][y][z]);
			}
		}
	}

	mHasChildren = false;
}

void Octree::Remove(Ball* ball, glm::vec3 pos)
{
	mBallNum--;

	if (mHasChildren && mBallNum < MIN_BALLS_PER_OCTREE)
	{
		DestroyChildren();
	}

	if (mHasChildren)
	{
		FileBall(ball, pos, false);
	}
	else
	{
		mBalls.erase(ball);
	}
}


void Octree::PotentialBallWallCollisions(std::vector<BallWallPair>& cs, Wall w, char coord, int dir)
{
	if (mHasChildren)
	{
		for (int dir2 = 0; dir2 < 2; dir2++)
		{
			for (int dir3 = 0; dir3 < 2; dir3++)
			{
				Octree* child = NULL;
				switch (coord)
				{
				case 'x':
					child = mChildren[dir][dir2][dir3];
					break;
				case 'y':
					child = mChildren[dir2][dir][dir3];
					break;
				case 'z':
					child = mChildren[dir2][dir3][dir];
					break;
				default:
					break;
				}

				child->PotentialBallWallCollisions(cs, w, coord, dir);
			}
		}
	}
	else
	{
		for (std::set<Ball*>::iterator it = mBalls.begin(); it != mBalls.end(); it++)
		{
			BallWallPair bwp;
			bwp.wall = w;
			bwp.ball = *it;
			cs.push_back(bwp);
		}
	}
}

void Octree::PotentialBallWallCollisions(std::vector<BallWallPair>& collisions)
{
	PotentialBallWallCollisions(collisions, WALL_LEFT, 'x', 0);
	PotentialBallWallCollisions(collisions, WALL_RIGHT, 'x', 1);
	PotentialBallWallCollisions(collisions, WALL_BOTTOM, 'y', 0);
	PotentialBallWallCollisions(collisions, WALL_TOP, 'y', 1);
	PotentialBallWallCollisions(collisions, WALL_FAR, 'z', 0);
	PotentialBallWallCollisions(collisions, WALL_NEAR, 'z', 1);
}

void Octree::PotentialBallBallCollisions(std::vector<BallPair>& collisions)
{
	if (mHasChildren)
	{
		for (int x = 0; x < 2; x++)
		{
			for (int y = 0; y < 2; y++)
			{
				for (int z = 0; z < 2; z++)
				{
					mChildren[x][y][z]->PotentialBallBallCollisions(collisions);
				}
			}
		}
	}
	else
	{
		for (std::set<Ball*>::iterator it = mBalls.begin(); it != mBalls.end(); it++)
		{
			for (std::set<Ball*>::iterator it2 = mBalls.begin(); it2 != mBalls.end(); it2++)
			{
				if (*it < *it2)
				{
					BallPair bp;
					bp.ball1 = *it;
					bp.ball2 = *it2;
					collisions.push_back(bp);
				}
			}

		}
	}
}

void PotentialBallBallCollisions(std::vector<BallPair>& potentialCollisions,
	std::vector<Ball*>& balls, Octree* octree)
{
	octree->PotentialBallBallCollisions(potentialCollisions);
}

void PotentialBallWallCollisions(std::vector<BallWallPair>& potentialCollisions,
	std::vector<Ball*>& balls, Octree* octree)
{
	octree->PotentialBallWallCollisions(potentialCollisions);
}

void MoveBalls(std::vector<Ball*>& balls, Octree* octree, float dt)
{
	for (int i = 0; i < balls.size(); i++)
	{
		Ball* ball = balls[i];
		glm::vec3 oldPosition = ball->position;
		ball->position += ball->velocity * dt;
		octree->MoveBall(ball, oldPosition);
	}
}

void ApplyGravity(std::vector<Ball*>& balls)
{
	for (int i = 0; i < balls.size(); i++)
	{
		Ball* ball = balls[i];
		ball->velocity -= glm::vec3(0, GRAVITY * TIME_BETWEEN_UPDATES, 0);
	}
}

bool TestBallBallCollision(Ball* b1, Ball* b2)
{
	float r = b1->radius + b2->radius;
	glm::vec3 delta = b1->position - b2->position;
	if (glm::dot(delta, delta) < r * r)
	{
		glm::vec3 netVelocity = b1->velocity - b2->velocity;
		glm::vec3 displacement = b1->position - b2->position;
		return glm::dot(netVelocity, displacement) < 0;
	}
	else
		return false;
}

void HandleBallBallCollisions(std::vector<Ball*>& balls, Octree* octree)
{
	std::vector<BallPair> bps;
	PotentialBallBallCollisions(bps, balls, octree);
	for (unsigned int i = 0; i < bps.size(); i++)
	{
		BallPair bp = bps[i];

		Ball* b1 = bp.ball1;
		Ball* b2 = bp.ball2;
		if (TestBallBallCollision(b1, b2))
		{
			glm::vec3 displacement = glm::normalize(b1->position - b2->position);
			b1->velocity -= displacement * glm::dot(b1->velocity, displacement) * 2.f;
			b2->velocity -= displacement * glm::dot(b2->velocity, displacement) * 2.f;
		}
	}
}

glm::vec3 WallDirection(Wall wall)
{
	switch (wall)
	{
	case WALL_LEFT:
		return glm::vec3(-1, 0, 0);
	case WALL_RIGHT:
		return glm::vec3(1, 0, 0);
	case WALL_FAR:
		return glm::vec3(0, 0, -1);
	case WALL_NEAR:
		return glm::vec3(0, 0, 1);
	case WALL_TOP:
		return glm::vec3(0, 1, 0);
	case WALL_BOTTOM:
		return glm::vec3(0, -1, 0);
	default:
		return glm::vec3(0, 0, 0);
	}
}

bool TestBallWallCollision(Ball* ball, Wall wall)
{
	glm::vec3 dir = WallDirection(wall);
	return glm::dot(ball->position, dir) + ball->radius > BOX_SIZE / 2 && glm::dot(ball->velocity, dir) > 0;
}

void HandleBallWallCollisions(std::vector<Ball*>& balls, Octree* octree)
{
	std::vector<BallWallPair> bwps;
	PotentialBallWallCollisions(bwps, balls, octree);
	for (unsigned int i = 0; i < bwps.size(); i++)
	{
		BallWallPair bwp = bwps[i];

		Ball* b = bwp.ball;
		Wall w = bwp.wall;
		if (TestBallWallCollision(b, w))
		{
			glm::vec3 dir = glm::normalize(WallDirection(w));
			b->velocity -= dir * glm::dot(b->velocity, dir) * 2.f;
		}
	}
}

void PerformUpdate(std::vector<Ball*>& balls, Octree* octree)
{
	ApplyGravity(balls);
	HandleBallBallCollisions(balls, octree);
	HandleBallWallCollisions(balls, octree);
}

void Advance(std::vector<Ball*>& balls, Octree* octree, float t, float& timeUntilUpdate)
{
	while (t > 0)
	{
		if (timeUntilUpdate <= t)
		{
			MoveBalls(balls, octree, timeUntilUpdate);
			PerformUpdate(balls, octree);
			t -= timeUntilUpdate;
			timeUntilUpdate = TIME_BETWEEN_UPDATES;
		}
		else
		{
			MoveBalls(balls, octree, t);
			timeUntilUpdate -= t;
			t = 0;
		}
	}
}


